package graphAlg.unioncollection;

import graphAlg.Graph;
import graphAlg.Node;

import java.util.*;

public class UnionFind {
    public  HashMap<Node,Node> fatherMap;
    public  HashMap<Node,Integer> sizeMap;

    public  void initUnionFind(Graph graph){
        fatherMap = new HashMap<>();
        sizeMap = new HashMap<>();
        HashMap<Integer, Node> nodes = graph.nodes;
        for (Node node : nodes.values()) {
            fatherMap.put(node,node);
            sizeMap.put(node,1);
        }
    }

    public boolean isSameSet(Node node1, Node node2){
        Node node1Father = getFather(node1);
        Node node2Father = getFather(node2);
        return node1Father == node2Father;

    }

    public  Node getFather(Node node){
        Stack<Node>  stack = new Stack<>();
        while(node != fatherMap.get(node)){
            stack.push(node);
            node = fatherMap.get(node);
        }
        while (!stack.isEmpty()){
            fatherMap.put(stack.pop(),node);
        }
        return node;
    }

    public void union(Node node1,Node node2){
        if(!isSameSet(node1,node2)){
            Node node1Father = getFather(node1);
            Node node2Father = getFather(node2);
            int size1 = sizeMap.get(node1Father);
            int size2 = sizeMap.get(node2Father);
            Node big = size1 >= size2 ? node1Father : node2Father;
            Node small = big == node1Father ? node2Father : node1Father;
            fatherMap.put(small,big);
            sizeMap.put(big,size1+size2);
            sizeMap.remove(small);
        }
    }
}
